import os
import json
import re
import config
import tqdm
from collections import OrderedDict
from util.get_model_instance import get_model_instance


def get_score(text):
    # 使用正则表达式提取一位小数
    match = re.search(r'分数为(\d+\.\d)分', text)
    if match:
        score = float(match.group(1))
        return score
    else:
        return 0


if __name__ == '__main__':
    qa_pairs = []
    model_name = config.MODEL_NAME
    rate_prompt = config.RATE_PROMPT
    batch_dir = '../mydata/'+model_name+'_generations_'+config.BATCH_DIR_SUFFIX+'/'
    output_file = 'machine_generated_instances.jsonl'
    instances_file_path = os.path.join(batch_dir, output_file)

    model_instance = get_model_instance(model_name)

    rate_output_file = 'qa_pairs_with_score.jsonl'
    rate_output_path = os.path.join(batch_dir, rate_output_file)

    # 统计行数
    with open(instances_file_path, 'r', encoding='utf-8') as file:
        total_num = sum(1 for line in file)
    progress_bar = tqdm.tqdm(total=total_num)


    existing_requests = {}
    if os.path.exists(rate_output_path):
        with open(rate_output_path, encoding='utf-8') as fin:
            for line in tqdm.tqdm(fin):
                try:
                    data = json.loads(line)
                    existing_requests[data["instruction"]] = data   # 这是个字典，key是instruction，value是整行data
                except:
                    pass
        print(f"Loaded {len(existing_requests)} existing requests")


    with open(instances_file_path, 'r', encoding='utf-8') as file:
        with open(rate_output_path, "w", encoding='utf-8') as fout:
            for line in file:
                data = json.loads(line)
                instruction =data['instruction']
                if instruction in existing_requests:
                    data = existing_requests[instruction]
                    data = OrderedDict(
                        (k, data[k]) for k in \
                        ["instruction", "output", "score"]
                    )
                    fout.write(json.dumps(data, ensure_ascii=False) + "\n")
                else:
                    output = data['output']
                    # 将instruction和output拼接成问答对, 进而拼接成prompt
                    qa_pair = f"Question: {instruction}\nAnswer: {output}"
                    prompt = rate_prompt + qa_pair
                    text = model_instance.get_raw_output(prompt)
                    # 后处理得分
                    score = get_score(text)
                    retries = 0
                    while score == 0 and retries <= 3:
                        score = get_score(text)
                    data['score'] = score
                    data = OrderedDict(
                            (k, data[k]) for k in \
                            ["instruction", "output",
                             "score"]
                        )
                    fout.write(json.dumps(data, ensure_ascii=False) + "\n")
                progress_bar.update(1)
